package com.webcat.service;

import com.webcat.config.openai.model.MyOpenAiChatModel;
import com.webcat.dify.kb.DifyKBClient;
import com.webcat.dify.kb.entity.Record;
import com.webcat.entity.SseResponse;
import com.webcat.entity.vo.DocumentVo;
import com.webcat.entity.vo.KnowledgeBaseVo;
import com.webcat.historymemory.PostgresChatMemory;
import com.webcat.utils.WordUtil;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;
import java.io.InputStream;
import java.util.List;
import java.util.UUID;

public abstract class AbstractChatService {

    @Resource(name = "myOpenAiChatModel")
    protected MyOpenAiChatModel openAiChatModel;

    @Resource
    protected ToolCallbackProvider toolCallbackProvider;

    @Resource
    protected DifyKBClient difyKBClient;

    @Resource
    private PostgresChatMemory postgresChatMemory;

    public String start() {
        return UUID.randomUUID().toString().toLowerCase();
    }

    protected ChatClient generateChatClient(String chatId, Boolean multiConversation) {
        return generateChatClient(chatId, multiConversation, true);
    }

    protected ChatClient generateChatClient(String chatId, Boolean multiConversation, boolean tool) {
        ChatClient.Builder builder = ChatClient.builder(this.openAiChatModel);
        if (tool) {
            if (toolCallbackProvider.getToolCallbacks().length > 0) {
                builder.defaultTools(toolCallbackProvider.getToolCallbacks());
            }
        }
        builder.defaultAdvisors(new SimpleLoggerAdvisor());
        ChatClient chatClient;
        if (multiConversation) {
            builder.defaultAdvisors(new MessageChatMemoryAdvisor(postgresChatMemory, chatId, 1000));
        }
        chatClient = builder.build();
        return chatClient;
    }

    protected SseResponse<Object> generateResponse(String chatId, ChatResponse response) {
        String finishReason = response.getResult().getMetadata().getFinishReason();
        if ("stop".equalsIgnoreCase(finishReason)) {
            return SseResponse.success(chatId, "end", "");
        } else {
            return SseResponse.success(chatId, "message", response.getResult().getOutput().getText());
        }
    }

    protected KnowledgeBaseVo generateKnowledgeBase(String datasetId, List<Record> records) {
        KnowledgeBaseVo knowledgeBaseVo = null;
        if (records != null && !records.isEmpty()) {
            knowledgeBaseVo = new KnowledgeBaseVo();
            knowledgeBaseVo.setDatasetId(datasetId);
            List<DocumentVo> documentVoList = records.stream().map(record -> {
                DocumentVo documentVo = new DocumentVo();
                documentVo.setId(record.getSegment().getDocument().getId());
                documentVo.setName(record.getSegment().getDocument().getName());
                documentVo.setContent(record.getSegment().getContent());
                return documentVo;
            }).toList();
            knowledgeBaseVo.setDocumentList(documentVoList);
        }
        return knowledgeBaseVo;
    }

    protected String generateRagContent(List<Record> records) {
        String ragContent = "";
        if (records != null && !records.isEmpty()) {
            StringBuilder messageBuilder = new StringBuilder();
            for (Record record: records) {
                String content = record.getSegment().getContent();
                messageBuilder.append(content).append("\n");
            }
            ragContent = messageBuilder.toString();
        }
        return ragContent;
    }

    protected Flux<SseResponse<Object>> streamChatWithFile(String chatId, String message, MultipartFile file, Boolean multiConversation) {

        try(InputStream in = file.getInputStream()) {
            String fileName = file.getOriginalFilename();
            int index = fileName.lastIndexOf(".");
            String extName = "";
            if (index != -1) {
                extName = fileName.substring(index+1);
            }

            String content = "";
            if ("doc".equalsIgnoreCase(extName) || "docx".equalsIgnoreCase(extName)) {
                content = WordUtil.parseWord(in);
            }
            String lastMessage = content + "\n" + message;
            Prompt prompt = new Prompt(lastMessage);

            Flux<ChatResponse> chatResponseFlux = generateChatClient(chatId, false).prompt(prompt).stream().chatResponse();

            return chatResponseFlux.map(response -> generateResponse(chatId, response));
        }
        catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
}